import os

import numpy as np
import torch
import torch.nn as nn
from PIL import Image
from sklearn.metrics.pairwise import cosine_similarity
from torchvision import models, transforms

from config.config import get_config

# 1. 加载预训练的 ResNet50 模型用于特征提取
model = models.resnet50(pretrained=True)
model = nn.Sequential(*list(model.children())[:-1])  # 去掉最后的全连接层
model.eval()

# 图像预处理
preprocess = transforms.Compose([
    transforms.Resize(512),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
features = []
filenames = []
root_path = get_config()['root']
class ImageTools:
    limit_load = -1
    # 2. 提取单张图像的特征向量
    def extract_features(self,image_path):
        img = Image.open(image_path).convert('RGB')
        img_t = preprocess(img)
        batch_t = torch.unsqueeze(img_t, 0)
        with torch.no_grad():
            feature = model(batch_t)
        return feature.numpy().flatten()

    def load_file(self,image_folder):
        if self.limit_load>0:
            if filenames.__len__() >= self.limit_load:
                return
        if os.path.exists(image_folder):
            files = os.listdir(image_folder)
            for file in files:
                file_dir = os.path.join(image_folder,file)
                if os.path.isdir(file_dir):
                    self.load_file(file_dir)
                elif file.endswith(('.png', '.jpg')):
                    relative_path = file_dir.replace(root_path, '')
                    try:
                        feature = self.extract_features(file_dir)
                        features.append(feature)
                        filenames.append(relative_path)
                        print(f"资源:\t{relative_path}\t加载中...")
                    except Exception as e:
                        print(e)
                        pass


    # 3. 构建图像数据库的特征库
    def build_feature_database(self,image_folder:str,max_limit:int=-1):
        self.limit_load=max_limit
        self.load_file(image_folder)
        return np.array(features), filenames


    # 4. 搜索最相似的图像
    def search_similar_images(self,query_image_path, feature, filename, top_k=5):
        query_feat = self.extract_features(query_image_path)
        query_feat = query_feat.reshape(1, -1)
        features_matrix = np.array(feature)

        # 计算余弦相似度
        similarities = cosine_similarity(query_feat, features_matrix)[0]

        # 获取最相似的 top_k 图像
        top_indices = np.argsort(similarities)[::-1][:top_k]

        results = [(filenames[i], similarities[i]) for i in top_indices]
        return results
